import os
import argparse
import numpy as np
import torch
from gan_model import Generator as VanillaGenerator
from wgan_model import Generator as WGANGenerator
from cgan_model import Generator as CGANGenerator
from wgan_gp_model import Generator as WGANGPGenerator

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--gan_type", type=str, default="vanilla", 
                       choices=["vanilla", "wgan", "cgan", "wgan_gp"], help="GAN类型")
    parser.add_argument("--n_samples", type=int, default=1000, help="每个类别生成的样本数")
    parser.add_argument("--latent_dim", type=int, default=100, help="潜在空间维度")
    parser.add_argument("--signal_length", type=int, default=1200, help="信号长度")
    parser.add_argument("--use_best_model", action="store_true", help="是否使用最佳模型")
    return parser.parse_args()

def generate_samples_for_class(class_name, args, device):
    """为特定类别生成样本"""
    # 类别映射
    class_mapping = {
        'normal': 0,
        'ball_07': 1, 'ball_14': 2, 'ball_21': 3,
        'inner_07': 4, 'inner_14': 5, 'inner_21': 6,
        'outer_07': 7, 'outer_14': 8, 'outer_21': 9
    }
    
    if class_name not in class_mapping:
        raise ValueError(f"未知的类别名称: {class_name}")
    
    class_idx = class_mapping[class_name]
    
    # 加载模型
    if args.gan_type == "cgan":
        model_dir = os.path.join("models", args.gan_type)
    else:
        model_dir = os.path.join("models", args.gan_type, class_name)
    
    # 选择使用最佳模型还是最终模型
    if args.use_best_model:
        model_path = os.path.join(model_dir, "generator_best.pth")
        if not os.path.exists(model_path):
            print(f"警告：找不到最佳模型 {model_path}，将使用最终模型")
            model_path = os.path.join(model_dir, "generator_final.pth")
    else:
        model_path = os.path.join(model_dir, "generator_final.pth")
    
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"找不到模型文件：{model_path}")
    
    print(f"使用模型：{model_path}")
    
    # 初始化生成器
    if args.gan_type == "vanilla":
        generator = VanillaGenerator(args.latent_dim, args.signal_length).to(device)
    elif args.gan_type == "wgan":
        generator = WGANGenerator(args.latent_dim, args.signal_length).to(device)
    elif args.gan_type == "cgan":
        generator = CGANGenerator(args.latent_dim, args.signal_length, n_classes=10).to(device)
    elif args.gan_type == "wgan_gp":
        generator = WGANGPGenerator(args.latent_dim, args.signal_length).to(device)
    
    # 加载模型权重
    try:
        state_dict = torch.load(model_path, map_location=device)
        generator.load_state_dict(state_dict)
    except Exception as e:
        print(f"加载模型权重时出错：{str(e)}")
        raise
    
    generator.eval()
    
    # 生成样本
    with torch.no_grad():
        z = torch.randn(args.n_samples, args.latent_dim).to(device)
        if args.gan_type == "cgan":
            # 为CGAN创建标签（所有样本都属于指定类别）
            labels = torch.full((args.n_samples,), class_idx, dtype=torch.long, device=device)
            samples = generator(z, labels)
        else:
            samples = generator(z)
    
    return samples.cpu().numpy()

def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 创建保存目录
    save_dir = os.path.join("cwru_gan", args.gan_type)
    if args.use_best_model:
        save_dir = os.path.join(save_dir, "best_model")
    os.makedirs(save_dir, exist_ok=True)
    
    # 类别列表
    class_names = [
        'normal',
        'ball_07', 'ball_14', 'ball_21',
        'inner_07', 'inner_14', 'inner_21',
        'outer_07', 'outer_14', 'outer_21'
    ]
    
    # 为每个类别生成样本
    for class_name in class_names:
        print(f"为 {class_name} 类别生成样本...")
        try:
            samples = generate_samples_for_class(class_name, args, device)
            save_path = os.path.join(save_dir, f"{class_name}.npy")
            np.save(save_path, samples)
            print(f"已保存 {len(samples)} 个样本到 {save_path}")
        except Exception as e:
            print(f"生成 {class_name} 类别的样本时出错：{str(e)}")

if __name__ == "__main__":
    main() 